#pragma once

#include <AVUIVisualTreeHelper.h>

namespace AVUI {

class TreeWalkPriority
{
public:
    enum Enum
    {
        LogicalTree,
        VisualTree
    };
};

class DescendentsWalkerBase
{
public:
#pragma region Internal Methods

    bool WasVisited(DependencyObject* pDO)
    {
        DependencyObject* pDependencyObject = pDO;
        while (pDependencyObject != m_pStartNode && pDependencyObject != NULL)
        {
            DependencyObject* pParent;
            if (is_type<FrameworkElement>(pDependencyObject))
            {
                FrameworkElement* pReference = (FrameworkElement*) pDependencyObject;;
                pParent = pReference->get_Parent();
                DependencyObject* pObj4 = VisualTreeHelper::GetParent(pReference);
                if (pObj4 != NULL && pParent != NULL && pObj4 != pParent)
                {
                    return m_nodes.Contains(pDependencyObject);
                }
                if (pObj4 == NULL)
                {
                    pDependencyObject = pParent;
                    continue;
                }
                pDependencyObject = pObj4;
                continue;
            }
            FrameworkContentElement* pElement2 = object_cast<FrameworkContentElement>(pDependencyObject);
            pParent = (pElement2 != NULL) ? pElement2->get_Parent() : NULL;
            pDependencyObject = pParent;
        }
        return pDependencyObject != NULL;
    }

#pragma region Internal Methods

#pragma region Internal Fields

    ItemStructList< TRefCountedPtr<DependencyObject> > m_nodes;
    TreeWalkPriority::Enum m_priority;
    int m_recursionDepth;
    TRefCountedPtr<DependencyObject> m_pStartNode;
    static const int MAX_TREE_DEPTH = 250;

#pragma endregion Internal Fields

protected:

#pragma region Protected Methods
    DescendentsWalkerBase(TreeWalkPriority::Enum priority) : m_pStartNode(NULL), m_priority(priority), m_recursionDepth(0)
    {
    }
#pragma endregion Protected Methods

};


template<class T>
class DescendentsWalker : public DescendentsWalkerBase
{
public:
    typedef bool (*VisitedCallback)(DependencyObject* pDO, T data);
#pragma region Internal Methods

    DescendentsWalker(TreeWalkPriority::Enum priority, VisitedCallback callback) : m_callback(callback), DescendentsWalkerBase(priority) { }
    DescendentsWalker(TreeWalkPriority::Enum priority, VisitedCallback callback, T data) : m_callback(callback), m_data(data), DescendentsWalkerBase(priority) { }

    void StartWalk(DependencyObject* pStartNode) { StartWalk(pStartNode, false); }
    virtual void StartWalk(DependencyObject* pStartNode, bool skipStartNode)
    {
        m_pStartNode = pStartNode;
        bool flag = true;
        if (!skipStartNode && (is_type<FrameworkElement>(m_pStartNode)) || is_type<FrameworkContentElement>(m_pStartNode))
        {
            flag = m_callback(m_pStartNode, m_data);
        }
        if (flag)
        {
            IterateChildren(m_pStartNode);
        }
    }

#pragma endregion Internal Methods

protected:

#pragma region Protected Methods

    virtual void _VisitNode(DependencyObject* pDO)
    {
        if (m_callback(pDO, m_data))
        {
            IterateChildren(pDO);
        }
    }

    T get_Data() { return m_data; }

#pragma endregion Protected Methods

private:

#pragma region Private Methods

    void IterateChildren(DependencyObject* pDO)
    {
        m_recursionDepth++;
        if (is_type<FrameworkElement>(pDO))
        {
            FrameworkElement* pFeParent = (FrameworkElement*) pDO;
            bool hasLogicalChildren = pFeParent->get_HasLogicalChildren();
            if (m_priority == TreeWalkPriority::VisualTree)
            {
                WalkFrameworkElementVisualThenLogicalChildren(pFeParent, hasLogicalChildren);
            }
            else if (m_priority == TreeWalkPriority::LogicalTree)
            {
                WalkFrameworkElementLogicalThenVisualChildren(pFeParent, hasLogicalChildren);
            }
        }
        else if (is_type<FrameworkContentElement>(pDO))
        {
            FrameworkContentElement* pFceParent = (FrameworkContentElement*)pDO;
            if (pFceParent->get_HasLogicalChildren())
            {
                WalkLogicalChildren(NULL, pFceParent, pFceParent->get_LogicalChildren());
            }
        }
        else
        {
            Visual* pVisual = object_cast<Visual>(pDO);
            if (pVisual != NULL)
            {
                WalkVisualChildren(pVisual);
            }
            else
            {
                Visual3D* pVisuald = object_cast<Visual3D>(pDO);
                if (pVisuald != NULL)
                {
                    WalkVisualChildren(pVisuald);
                }
            }
        }
        m_recursionDepth--;
    }

    void VisitNode(DependencyObject* pDO)
    {
        if (m_recursionDepth > 250)
        {
            LibraryCriticalError();
        }
        if (is_type<FrameworkElement>(pDO))
        {
            VisitNode((FrameworkElement*)pDO);
        }
        else if (is_type<FrameworkContentElement>(pDO))
        {
            _VisitNode(pDO);
        }
        else
        {
            IterateChildren(pDO);
        }

    }

    void VisitNode(FrameworkElement* pFe)
    {
        if (m_recursionDepth > 250)
        {
            LibraryCriticalError();
        }
        int index = m_nodes.IndexOf(pFe);
        if (index != -1)
        {
            m_nodes.RemoveAt(index);
        }
        else
        {
            DependencyObject* pParent = VisualTreeHelper::GetParent(pFe);
            DependencyObject* pObj3 = pFe->get_Parent();
            if (pParent != NULL && pObj3 != NULL && pParent != pObj3)
            {
                m_nodes.AddItem(pFe);
            }
            _VisitNode(pFe);
        }
    }

    void WalkFrameworkElementLogicalThenVisualChildren(FrameworkElement* pFeParent, bool hasLogicalChildren)
    {
        if (hasLogicalChildren)
        {
            WalkLogicalChildren(pFeParent, NULL, pFeParent->get_LogicalChildren());
        }

        pFeParent->set_IsVisualChildrenIterationInProgress(true);

        AVUI_TRY
        {
            int internalVisualChildrenCount = pFeParent->get_InternalVisualChildrenCount();
            for (int i = 0; i < internalVisualChildrenCount; i++)
            {
                Visual* pVisualChild = pFeParent->InternalGetVisualChild(i);
                if (pVisualChild != NULL && is_type<FrameworkElement>(pVisualChild) && VisualTreeHelper::GetParent(pVisualChild) != ((FrameworkElement*) pVisualChild)->get_Parent())
                {
                    VisitNode(pVisualChild);
                }
            }
        }
        AVUI_FINALLY
        {
            pFeParent->set_IsVisualChildrenIterationInProgress(false);
        }
        TRefCountedPtr<ItemStructList<TRefCountedPtr<Popup> >  > pList = Popup::RegisteredPopupsField.GetValue(pFeParent);
        if (pList != NULL)
        {
            for(int i = 0; i < pList->get_Count(); i++)
            {
                VisitNode((FrameworkElement*)(*pList)[i]);
            }
        }
    }

    void WalkFrameworkElementVisualThenLogicalChildren(FrameworkElement* pFeParent, bool hasLogicalChildren)
    {
        WalkVisualChildren(pFeParent);
        TRefCountedPtr<ItemStructList<TRefCountedPtr<Popup> >  > pList = Popup::RegisteredPopupsField.GetValue(pFeParent);
        if (pList != NULL)
        {
            for(int i = 0; i < pList->get_Count(); i++)
            {
                VisitNode((FrameworkElement*)(*pList)[i]);
            }
        }

        pFeParent->set_IsLogicalChildrenIterationInProgress(true);
        AVUI_TRY
        {
            if (hasLogicalChildren)
            {
                TRefCountedPtr<Object> pLogicalChildrenObject = pFeParent->get_LogicalChildren();
                IEnumerator* pLogicalChildren = interface_cast<IEnumerator>(pLogicalChildrenObject);
                if (pLogicalChildren != NULL)
                {
                    while (pLogicalChildren->MoveNext())
                    {
                        Object* pCurrent = pLogicalChildren->get_Current();
                        FrameworkElement* pReference = object_cast<FrameworkElement>(pCurrent);
                        if (pReference != NULL)
                        {
                            if (VisualTreeHelper::GetParent(pReference) != pReference->get_Parent())
                            {
                                VisitNode(pReference);
                            }
                        }
                        else
                        {
                            FrameworkContentElement* pFCE = object_cast<FrameworkContentElement>(pCurrent);
                            if (pFCE != NULL)
                            {
                                VisitNode(pFCE);
                            }
                        }
                    }
                }
            }
        }
        AVUI_FINALLY
        {
            pFeParent->set_IsLogicalChildrenIterationInProgress(false);
        }

    }

    void WalkLogicalChildren(FrameworkElement* pFeParent, FrameworkContentElement* pFceParent, Object* pLogicalChildrenObject)
    {
        IEnumerator* pLogicalChildren = interface_cast<IEnumerator>(pLogicalChildrenObject);
        if (pFeParent != NULL)
        {
            pFeParent->set_IsLogicalChildrenIterationInProgress(true);
        }
        else
        {
            pFceParent->set_IsLogicalChildrenIterationInProgress(true);
        }
        AVUI_TRY
        {
            if (pLogicalChildren != NULL)
            {
                while (pLogicalChildren->MoveNext())
                {
                    DependencyObject* pCurrent = object_cast<DependencyObject>(pLogicalChildren->get_Current());
                    if (pCurrent != NULL)
                    {
                        VisitNode(pCurrent);
                    }
                }
            }
        }
        AVUI_FINALLY
        {
            if (pFeParent != NULL)
            {
                pFeParent->set_IsLogicalChildrenIterationInProgress(false);
            }
            else
            {
                pFceParent->set_IsLogicalChildrenIterationInProgress(false);
            }
        }
    }

    void WalkVisualChildren(Visual3D* pV)
    {
        pV->set_IsVisualChildrenIterationInProgress(true);
        AVUI_TRY
        {
            int num = pV->get_InternalVisual2DOr3DChildrenCount();
            for (int i = 0; i < num; i++)
            {
                DependencyObject* pDO = pV->InternalGet2DOr3DVisualChild(i);
                if (pDO != NULL)
                {
                    VisitNode(pDO);
                }
            }
        }
        AVUI_FINALLY
        {
            pV->set_IsVisualChildrenIterationInProgress(false);
        }
    }

    void WalkVisualChildren(Visual* pV)
    {
        pV->set_IsVisualChildrenIterationInProgress(true);
        AVUI_TRY
        {
            int num = pV->get_InternalVisual2DOr3DChildrenCount();
            for (int i = 0; i < num; i++)
            {
                DependencyObject* pDO = pV->InternalGet2DOr3DVisualChild(i);
                if (pDO != NULL)
                {
                    VisitNode(pDO);
                }
            }
        }
        AVUI_FINALLY
        {
            pV->set_IsVisualChildrenIterationInProgress(false);
        }
    }

#pragma endregion Private Methods

    VisitedCallback m_callback;
    T m_data;
};


template<class T>
class PrePostDescendentsWalker : public DescendentsWalker<T>
{
public:
#pragma region Internal Methods

    PrePostDescendentsWalker(TreeWalkPriority::Enum priority, VisitedCallback preCallback, VisitedCallback postCallback, T data) : m_postCallback(postCallback), DescendentsWalker<T>(priority, preCallback, data) { }

    void StartWalk(DependencyObject* pStartNode) { StartWalk(pStartNode, false); }
    virtual void StartWalk(DependencyObject* pStartNode, bool skipStartNode) AVUI_OVERRIDE
    {
        AVUI_TRY
        {
            DescendentsWalker<T>::StartWalk(pStartNode, skipStartNode);
        }
        AVUI_FINALLY
        {
            if (!skipStartNode && m_postCallback != NULL && (is_type<FrameworkElement>(pStartNode) || is_type<FrameworkContentElement>(pStartNode)))
            {
                m_postCallback(pStartNode, get_Data());
            }
        }

    }

#pragma endregion Internal Methods

protected:

#pragma region Protected Methods

    virtual void _VisitNode(DependencyObject* pDO) AVUI_OVERRIDE
    {
        AVUI_TRY
        {
            DescendentsWalker<T>::_VisitNode(pDO);
        }
        AVUI_FINALLY
        {
            if (m_postCallback != NULL)
            {
                m_postCallback(pDO, get_Data());
            }
        }
    }

#pragma endregion Protected Methods

private:

#pragma region Private Fields

    VisitedCallback m_postCallback;

#pragma endregion Private Fields
};


};


